from langchain.chains import create_sql_query_chain
from langchain.prompts.prompt import PromptTemplate
from langchain import SQLDatabase
from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool

SUPPORT_DATABASE = ["mysql"]


DATABASE_ADDR = {
    "mysql": "mysql+pymysql://{}:{}@{}:{}/{}"
}


# 自定义prompt
_DEFAULT_TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Use the following format:

Question: "Question here"
SQLQuery: "SQL Query to run"
SQLResult: "Result of the SQLQuery"
Answer: "Final answer here"

Only use the following tables:

{table_info}

Question: {input}"""

PROMPT = PromptTemplate(
    input_variables=["input", "table_info", "dialect"], template=_DEFAULT_TEMPLATE
)


class SQLAgent:
    def __init__(self, type, user, password, ip, port, database, llm):
        self.type = type
        self.user = user
        self.password = password
        self.ip = ip
        self.port = port
        self.database = database
        self.llm = llm
        self.addr = DATABASE_ADDR[self.type].format(self.user, self.password, self.ip, self.port, self.database)        
        self.db = SQLDatabase.from_uri(self.addr)     
        self.chain = create_sql_query_chain(self.llm, self.db)

    def formatDatabaseConnect(self):
        return self.db
        
    def getSQLDatabaseChain(self):
        return self.chain
    
    def textToSQL(self, question: str):
        sql = self.chain.invoke({"question": question})
        sql = sql.split("SQLResult")[0]
        sql = sql.split(";")[0]
        sql += ";"
        return sql
        
    def sqlQuery(self, sql: str):
        return self.db.run(sql, include_columns=True)
    
    def ask(self, question: str):
        pass
    